{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from tqdm import tqdm\n",
    "import pickle\n",
    "\n",
    "import torch\n",
    "\n",
    "import hydra\n",
    "from hydra import initialize, compose\n",
    "\n",
    "from home_robot.mapping.voxel.feature_voxel import ConceptFusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from home_robot.datasets.scannet import ScanNetDataset, ReferIt3dDataConfig, ScanReferDataConfig\n",
    "data = ScanNetDataset(\n",
    "    root_dir = '/private/home/ssax/home-robot/src/home_robot/home_robot/datasets/scannet/data',\n",
    "    frame_skip = 180,\n",
    "    n_classes = 50,\n",
    "    referit3d_config = ReferIt3dDataConfig(),\n",
    "    scanrefer_config = ScanReferDataConfig(),\n",
    ")\n",
    "\n",
    "# Load specific scene\n",
    "# scene0192_00 -- small scene\n",
    "# 'scene0000_00' -- large scene\n",
    "idx = data.scene_list.index(\"scene0192_00\") #'scene0000_00'\n",
    "# idx = 0\n",
    "print(f\"Loaded images of (h: {data.height}, w: {data.width}) - resized from ({data.DEFAULT_HEIGHT},{data.DEFAULT_WIDTH})\")\n",
    "scene_obs = data.__getitem__(idx, show_progress=True)\n",
    "torch.autograd.set_grad_enabled(False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load parameters\n",
    "hydra.core.global_hydra.GlobalHydra.instance().clear()\n",
    "\n",
    "# hydra.initialize(config_path='configs', config_name='cf_modular')\n",
    "\n",
    "with initialize(version_base='1.3', config_path=\"configs/model\"):\n",
    "    cfg = compose(config_name=\"conceptfusion\", overrides=[])\n",
    "\n",
    "# Initialize ConceptFusion only once\n",
    "concept_fusion = hydra.utils.instantiate(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "concept_fusion.clear()\n",
    "\n",
    "\n",
    "# SETUP EVAL\n",
    "class_id_to_class_names = dict(\n",
    "    zip(\n",
    "        data.METAINFO[\"CLASS_IDS\"],  # IDs [1, 3, 4, 5, ..., 65]\n",
    "        data.METAINFO[\"CLASS_NAMES\"],  # [wall, floor, cabinet, ...]\n",
    "    )\n",
    ")\n",
    "\n",
    "concept_fusion.set_vocabulary(class_id_to_class_names)\n",
    "\n",
    "keys = [\n",
    "    \"images\",\n",
    "    \"depths\",\n",
    "    \"poses\",\n",
    "    \"intrinsics\",\n",
    "    \"boxes_aligned\",\n",
    "    \"box_classes\",\n",
    "]\n",
    "\n",
    "gt_bounds, gt_classes, pred_bounds, pred_classes, pred_scores = [], [], [], [], []\n",
    "# Move to device\n",
    "for k in keys:\n",
    "    scene_obs[k] = scene_obs[k].to(concept_fusion.device)\n",
    "\n",
    "# Eval each scene and move to CPU\n",
    "queries = {\n",
    "    int(clas): class_id_to_class_names[int(clas)]\n",
    "    for clas in scene_obs[\"box_classes\"].unique()\n",
    "}\n",
    "\n",
    "# concept_fusion.build_scene(scene_obs)\n",
    "# concept_fusion.show_point_cloud_pytorch3d()\n",
    "\n",
    "instances_dict = concept_fusion.build_scene_and_get_instances_for_queries(scene_obs, queries.values())\n",
    "concept_fusion.show_point_cloud_pytorch3d(instances_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pc_xyz, pc_feat, _, pc_rgb = concept_fusion.voxel_map.get_pointcloud()\n",
    "inds, similarity = concept_fusion.query_similarity(\"chair\", pc_feat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch3d.structures import Pointclouds\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs\n",
    "from home_robot.utils.bboxes_3d import BBoxes3D\n",
    "\n",
    "from home_robot.utils.bboxes_3d_plotly import plot_scene_with_bboxes\n",
    "from home_robot.utils.data_tools.dict import update\n",
    "import matplotlib\n",
    "traces = {}\n",
    "cmap = matplotlib.colormaps[\"jet\"]\n",
    "similarity_colormap = cmap(similarity.detach().cpu().numpy())[:, :3]\n",
    "\n",
    "map_colors = 0.5 * pc_rgb.cpu() / 255. + 0.5 * torch.tensor(similarity_colormap)\n",
    "map_colors = torch.tensor(similarity_colormap)\n",
    "traces[\"Points\"] = Pointclouds(points=[pc_xyz.cpu()], features=[map_colors])\n",
    "\n",
    "_default_plot_args = dict(\n",
    "            xaxis={\"backgroundcolor\": \"rgb(200, 200, 230)\"},\n",
    "            yaxis={\"backgroundcolor\": \"rgb(230, 200, 200)\"},\n",
    "            zaxis={\"backgroundcolor\": \"rgb(200, 230, 200)\"},\n",
    "            axis_args=AxisArgs(showgrid=True),\n",
    "            pointcloud_marker_size=3,\n",
    "            pointcloud_max_points=500_000,\n",
    "            boxes_plot_together=False,\n",
    "            boxes_wireframe_width=3,\n",
    ")\n",
    "fig = plot_scene_with_bboxes(\n",
    "            plots={f\"Conceptfusion Pointcloud\": traces},\n",
    "            **update(_default_plot_args, _default_plot_args),\n",
    ")\n",
    "fig.update_layout(\n",
    "            height=800,\n",
    "            width=1600,\n",
    ")\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instances_dict = concept_fusion.get_instances_for_queries(queries.values())\n",
    "concept_fusion.show_point_cloud_pytorch3d(instances_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Block to test hyperparams.\n",
    "\n",
    "# remove \"otherfurniture\" from queries to test performance without it\n",
    "# All furniture classes have high similarity to \"otherfurniture\" so it is not a good query\n",
    "queries.pop(39)\n",
    "\n",
    "concept_fusion.dbscan_params.epsilon = 0.2\n",
    "concept_fusion.dbscan_params.min_samples = 50\n",
    "concept_fusion.similarity_params.similarity_thresh = 0.2\n",
    "instances_dict = concept_fusion.text_queries(queries=queries.values())\n",
    "concept_fusion.show_point_cloud_pytorch3d(instances_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize how the ground truth boxes look like\n",
    "import open3d as o3d\n",
    "import numpy as np\n",
    "\n",
    "pcd = o3d.io.read_point_cloud(\"/srv/flash1/kyadav32/datasets/ovmm/scannet/scannet_scene0011_00_gt.ply\")\n",
    "\n",
    "from pytorch3d.structures import Pointclouds\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs\n",
    "from home_robot.utils.bboxes_3d import BBoxes3D\n",
    "\n",
    "from home_robot.utils.bboxes_3d_plotly import plot_scene_with_bboxes\n",
    "from home_robot.utils.data_tools.dict import update\n",
    "\n",
    "from utils import COLOR_LIST\n",
    "\n",
    "traces = {}\n",
    "\n",
    "pc_xyz, pc_rgb = pcd.points, pcd.colors\n",
    "pc_xyz = torch.tensor(np.asarray(pc_xyz), device=\"cuda\")\n",
    "pc_rgb = torch.tensor(np.asarray(pc_rgb), device=\"cuda\")\n",
    "\n",
    "traces[\"Points\"] = Pointclouds(points=[pc_xyz], features=[pc_rgb*255])\n",
    "\n",
    "\n",
    "box_classes = dataset_0[0]['box_classes']\n",
    "box_bounds = dataset_0[0]['boxes_aligned']\n",
    "\n",
    "bounds, names, colors = {}, {}, {}\n",
    "for class_id, bound in zip(box_classes, box_bounds):\n",
    "    \n",
    "    class_name = class_id_to_class_names[class_id.item()]\n",
    "\n",
    "    if class_name not in bounds:\n",
    "        bounds[class_name] = []\n",
    "        names[class_name] = []\n",
    "        colors[class_name] = []\n",
    "    bounds[class_name].append(bound)\n",
    "    names[class_name].append(torch.tensor(class_id, device='cuda'))\n",
    "    colors[class_name].append(torch.tensor(COLOR_LIST[class_id % len(COLOR_LIST)], device='cuda'))\n",
    "for class_name in box_classes.keys():\n",
    "    detected_boxes = BBoxes3D(\n",
    "        bounds=[torch.stack(bounds, dim=0)],\n",
    "        features=[torch.stack(colors, dim=0)],\n",
    "        names=[torch.stack(names, dim=0).unsqueeze(-1)],\n",
    "    )\n",
    "    traces[class_name + \"_bbox\"] = detected_boxes\n",
    "\n",
    "\n",
    "_default_plot_args = dict(\n",
    "    xaxis={\"backgroundcolor\": \"rgb(200, 200, 230)\"},\n",
    "    yaxis={\"backgroundcolor\": \"rgb(230, 200, 200)\"},\n",
    "    zaxis={\"backgroundcolor\": \"rgb(200, 230, 200)\"},\n",
    "    axis_args=AxisArgs(showgrid=True),\n",
    "    pointcloud_marker_size=3,\n",
    "    pointcloud_max_points=500_000,\n",
    "    boxes_plot_together=False,\n",
    "    boxes_wireframe_width=3,\n",
    ")\n",
    "fig = plot_scene_with_bboxes(\n",
    "    plots={f\"Conceptfusion Pointcloud\": traces},\n",
    "    **update(_default_plot_args, {}),\n",
    ")\n",
    "fig.update_layout(\n",
    "    height=800,\n",
    "    width=1600,\n",
    ")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
